# coding=utf-8

import paramiko
import time


class SSHRequest:

    def __init__(self, hostname, port, username, password, timeout, command):
        """
        :param hostname: 主机名
        :param port:     端口
        :param username: 用户
        :param password: 密码
        :param timeout:  tcp连接超时时间 单位ms
        :param command:  shell命令
        """
        self.hostname = hostname
        self.port = port
        self.username = username
        self.password = password
        self.timeout = timeout / 1000  # ms to s
        self.command = command
        # 请求与响应
        self.response_headers = None  # type: dict
        self.response_body = None  # type: dict
        self.request_headers = None  # type: dict
        self.request_body = None  # type: dict
        self.stdout = ''
        self.stderr = ''
        # 请求响应耗时 ms
        self.elapsed_time = 0

    def exec_command(self):
        """
        :return: stdin, stdout, stderr
        """
        # 打开一个Channel并执行命令
        return self.client.exec_command(self.command)  # stdout 为正确输出，stderr为错误输出，同时是有1个变量有值

    def send(self):
        with self:
            # send request
            _start_clock = time.time()
            stdin, stdout, stderr = self.exec_command()
            # recv response
            _end_clock = time.time()
            self.elapsed_time = int((_end_clock - _start_clock) * 1000) + 1
            self.stdout = stdout.read().decode('utf-8')
            self.stderr = stderr.read().decode('utf-8')
        self._handle_response()

    def close(self):
        self.client.close()

    def __enter__(self):
        # 实例化SSHClient
        self.client = paramiko.SSHClient()
        # 自动添加策略，保存服务器的主机名和密钥信息，如果不添加，那么不再本地know_hosts文件中记录的主机将无法连接
        # 报错：paramiko.ssh_exception.SSHException: Server '127.0.0.1(Server地址)' not found in known_hosts
        self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        # 连接SSH服务端，以用户名和密码进行认证
        self.client.connect(hostname=self.hostname, port=self.port, username=self.username, password=self.password,
                            timeout=self.timeout)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def _handle_response(self):
        """处理应答解析拿到 应答体 应答头 请求体 请求头"""
        self._handle_response_headers()
        self._handle_request_body()
        self._handle_request_headers()
        self._handle_response_body()

    def _handle_response_headers(self):
        self.response_headers = {
            'hostname': self.hostname,
            'port': self.port,
            'username': self.username,
            'password': self.password,
            'timeout': self.timeout,
            'command': self.command,
            'elapsed_time': self.elapsed_time,
        }

    def _handle_request_body(self):
        self.request_body = {
            'command': self.command,
        }

    def _handle_request_headers(self):
        self.request_headers = {
            'hostname': self.hostname,
            'port': self.port,
            'username': self.username,
            'password': self.password,
            'timeout': self.timeout,
            'command': self.command,
        }

    def _handle_response_body(self):
        self.response_body = {
            'stdout': self.stdout,
            'stderr': self.stderr,
        }


def make_request(hostname, port, username, password, timeout, command):
    """
    :param hostname: 主机名
    :param port:     端口
    :param username: 用户
    :param password: 密码
    :param timeout:  tcp连接超时时间 单位ms
    :param command:  shell命令
    """
    return SSHRequest(hostname=hostname, port=port, username=username, password=password, timeout=int(timeout),
                      command=command)


if __name__ == '__main__':
    pass
